from typing import Annotated
from typing_extensions import TypedDict

from langchain_core.messages import BaseMessage
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.checkpoint.postgres import PostgresSaver

from langchain_openai import ChatOpenAI
from langchain_core.tools import tool


# 1. 定义状态
class State(TypedDict):
    messages: Annotated[list, add_messages]


# 2. 定义工具
@tool
def calculate(expression: str) -> str:
    """🧮 数学计算工具"""
    print("calculate 工具被调用, 进行计算")
    try:
        allowed_chars = set('0123456789+-*/().,e ')
        if not all(c in allowed_chars for c in expression):
            return "错误：表达式包含不允许的字符"
        result = eval(expression)
        return f"计算结果：{expression} = {result}"
    except Exception as e:
        return f"计算错误：{str(e)}"


@tool
def get_weather(location: str) -> str:
    """🌤️ 天气查询工具"""
    print("get_weather 工具被调用, 进行天气查询")
    info = f"{location} 今天天气晴天☀️, 温度26°C, 空气质量良好"
    return info


# 3. 初始化模型
llm = ChatOpenAI(
    model_name="deepseek-chat",
    openai_api_key="sk-c2acb26542994445a95052fbcae9cd02",
    openai_api_base="https://api.deepseek.com/v1",
)

tools = [calculate, get_weather]
llm_with_tools = llm.bind_tools(tools)


# 4. 系统提示模板
SYSTEM_PROMPT_TEMPLATE = """你是一个智能助手，可以进行自然对话并在需要时使用专业工具来提供准确信息。

🛠️ 可用工具：
{tools_info}

📋 工具调用规则：
- 仔细分析用户意图，判断是否需要使用工具
- 如果需要使用工具，严格按照格式：[TOOL_CALL]工具名称("参数")[/TOOL_CALL]
- 参数必须用双引号包围
- 对于普通对话、知识问答、闲聊等，直接回答，无需使用工具

💡 使用示例：
- 天气查询："北京今天天气怎么样？" → [TOOL_CALL]get_weather("北京")[/TOOL_CALL]
- 数学计算："帮我计算 25 * 4 + 10" → [TOOL_CALL]calculate("25 * 4 + 10")[/TOOL_CALL]
- 普通对话："你好" → 你好！我是你的AI智能助手，有什么可以帮助你的吗？

🎯 重要提醒：
- 准确识别用户需求，避免不必要的工具调用
- 工具调用格式必须严格正确
- 始终保持友好、专业的对话态度
"""


# 5. 构建图
def build_graph(checkpointer):
    def chatbot(state: State):
        # 把 SYSTEM_PROMPT_TEMPLATE 注入为 system 消息
        system_msg = (
            "system",
            SYSTEM_PROMPT_TEMPLATE.format(
                tools_info="\n".join(
                    [f"- {t.name}: {t.description}" for t in tools]
                )
            ),
        )
        messages = [system_msg] + state["messages"]
        return {"messages": [llm_with_tools.invoke(messages)]}

    tool_node = ToolNode(tools=tools)

    graph_builder = StateGraph(State)
    graph_builder.add_node("chatbot", chatbot)
    graph_builder.add_node("tools", tool_node)

    graph_builder.add_conditional_edges("chatbot", tools_condition)
    graph_builder.add_edge("tools", "chatbot")
    graph_builder.set_entry_point("chatbot")

    return graph_builder.compile(checkpointer=checkpointer)


# 6. 运行交互
def run_chatbot(graph):
    print("Chatbot started! 输入 'quit' 退出。")
    while True:
        try:
            user_input = input("User: ")
            if user_input.lower() in ["quit", "exit", "q"]:
                print("Goodbye!")
                break
            for chunk in graph.stream(
                {"messages": [("user", user_input)]},
                config={"configurable": {"thread_id": 1}},
                stream_mode="values",
            ):
                print("AI:", chunk["messages"][-1].content)

        except (KeyboardInterrupt, EOFError):
            print("\nGoodbye!")
            break
        except Exception as e:
            import traceback
            print(f"An error occurred: {e}")
            traceback.print_exc()


# 7. 主函数
def main():
    DB_URI = "postgresql://postgres:123456@192.168.31.193:5433/llm?sslmode=disable"
    with PostgresSaver.from_conn_string(DB_URI) as checkpointer:
        graph = build_graph(checkpointer)
        run_chatbot(graph)


if __name__ == "__main__":
    main()
